# Copyright (c) 2019 Presto Labs Pte. Ltd.
# Author: taekwon(originally from huobi_futures/order_gateway.py)

import collections
import enum
import functools
import typing

import requests
import recordclass
import tornado.gen

import coin.flow.topic_map as topic_map
import coin.flow.subscriber as flow_subscriber
import coin.proto.coin_market_enums_pb2 as coin_enum

from coin.base.config import Config
from coin.base.timestamp import get_timestamp
from coin.exchange.base.order_fill_util import FillManager
from coin.exchange.base.order_gateway_ready import PositionOrderGatewayReady
from coin.exchange.base.rest.task_scheduler import Scheduler, wrap_async_task
from coin.exchange.base.rest_order_gateway_base import PrivateExchangeMessage
from coin.exchange.deribit_v1.kr_rest.currency import DeribitCurrency
from coin.exchange.deribit_v1.kr_rest.futures_product import DeribitFuturesProduct
from coin.exchange.deribit_v2.rest.private_client import DeribitFuturesPrivateParser
from coin.exchange.util.rest_rate_limit import RestRateLimiter
from coin.proto.coin_order_gateway_pb2 import OrderGatewayConfig, OrderEvent
from coin.strategy.hk_hard_arb.og_util import print_working_orders

from coin.proto.coin_query_pb2 import (
    AccountBalance,
    AccountPosition,
    ProductPosition,
)
from coin.exchange.deribit_v2.rest.native_private_client import DeribitFuturesNativePrivateClient
from coin.exchange.deribit_v2.ws.ws_private_client import DeribitFuturesWsPrivateClient

from coin.exchange.deribit_v2.rest.native_private_client import (
    DeribitFuturesDirectionType,
    DeribitFuturesOffsetType,
    DeribitFuturesOrderPriceType,
    DeribitFuturesStatusType,
)
from coin.exchange.base.order_gateway import (
    BaseOrderGatewayInfoMutator,
    Order,
    OrderGatewayStatus,
    OrderSide,
    OrderType,
)
from coin.exchange.base.rest_order_gateway_base import (
    RestOrderGateway,
    create_order_internal,
    exchange_message_handler,
)


# TODO(taekwon): consider various futures.
def translate_order_side(deribit_order_side: DeribitFuturesDirectionType,
                         deribit_order_type: DeribitFuturesOrderPriceType):
  order_side = OrderSide.UNKNOWN
  if deribit_order_side == DeribitFuturesDirectionType.BUY:
    order_side = OrderSide.BUY
  elif deribit_order_side == DeribitFuturesDirectionType.SELL:
    order_side = OrderSide.SELL
  order_type = OrderType.UNKNOWN
  if deribit_order_type == DeribitFuturesOrderPriceType.LIMIT:
    order_type = OrderType.LIMIT
  if deribit_order_type == DeribitFuturesOrderPriceType.MARKET:
    order_type = OrderType.MARKET

  return order_side, order_type


def create_order_from_open_order_message(message, timestamp, order_id=None):
  """
  this is for deribit.
  {
    "jsonrpc": "2.0",
    "id": 4316,
    "result": {
      "time_in_force": "good_til_cancelled",
      "reduce_only": false,
      "profit_loss": 0.051134,
      "price": 118.94,
      "post_only": false,
      "order_type": "limit",
      "order_state": "filled",
      "order_id": "ETH-331562",
      "max_show": 37,
      "last_update_timestamp": 1550219810944,
      "label": "",
      "is_liquidation": false,
      "instrument_name": "ETH-PERPETUAL",
      "filled_amount": 37,
      "direction": "sell",
      "creation_timestamp": 1550219749176,
      "commission": 0.000031,
      "average_price": 118.94,
      "api": false,
      "amount": 37
    }
  }

  """
  price_orig = float(message['price'])
  qty_orig = float(message['amount'])
  product = DeribitFuturesProduct.FromStrNativeProduct(message['instrument_name'])
  exchange_order_id = convert_to_exchange_order_id(product, message['order_id'])
  qty_accum_fill = float(message['filled_amount'])
  if message['average_price']:
    price_avg_fill = float(message['average_price'])
  else:
    price_avg_fill = None
  fee = float(message['commission'])
  status = DeribitFuturesStatusType(message['order_state'])
  order_side, order_type = translate_order_side(DeribitFuturesDirectionType(message['direction']),
                                                DeribitFuturesOrderPriceType(message['order_type']))

  fill = FillManager(
      price_orig=price_orig,
      qty_orig=qty_orig,
      price_last_fill=price_avg_fill,
      qty_last_fill=qty_accum_fill,
      price_avg_fill=price_avg_fill,
      qty_accum_fill=qty_accum_fill,
      timestamp_last_fill=timestamp,
  )

  order = Order(order_id=order_id,
                product=product,
                price=price_orig,
                qty=qty_orig,
                order_side=order_side,
                order_type=order_type,
                tag=None,
                accepted=True,
                cancel_sent=False,
                internal=create_order_internal(
                    timestamp=timestamp,
                    exchange_order_id=exchange_order_id,
                    fill=fill,
                    fee=fee,
                    status=status,
                ))
  return order


# `real_exchange_order_id` refers to the order_id chosen by the exchange.
# It is a very simple number, like 14, and looks like unique for one particular
# product and user account. Since we assume exchange order id is a unique id
# for all orders of one account. So we need to convert
# real_exchange_order_id to exchange_order_id.
def get_real_exchange_order_id(order):
  if order.internal.exchange_order_id is not None:
    return order.internal.exchange_order_id.split('_')[1]
  else:
    return None


def convert_to_exchange_order_id(product, real_exchange_order_id):
  return '%s_%s' % (product.symbol, real_exchange_order_id)


FuturesPosition = recordclass.recordclass('FuturesPosition',
                                          ['long', 'short', 'long_avg_price', 'short_avg_price'])


def create_futures_positions(long=0., short=0., long_avg_price=None, short_avg_price=None):
  return FuturesPosition(long=long,
                         short=short,
                         long_avg_price=long_avg_price,
                         short_avg_price=short_avg_price)


class DeribitFuturesOrderGatewayInfo(BaseOrderGatewayInfoMutator):
  CurrencyType = DeribitCurrency
  ProductType = DeribitFuturesProduct
  BalanceType = float
  PositionType = float

  def __init__(self,
               currencies: typing.List[DeribitCurrency],
               products: typing.List[DeribitFuturesProduct]):
    super().__init__(currencies, products)
    self._unrealized_pnl = collections.defaultdict(float)
    self._avg_entry_price = collections.defaultdict(create_futures_positions)

  def update_balance(self, message):
    msg = DeribitFuturesPrivateParser.parse_balance(message)

    for balance in msg.each_balance:
      currency = DeribitCurrency.FromStrNativeCurrency(balance.currency_native)
      if currency in self.currencies:
        self.set_balance(currency, balance.total, balance)

  def update_position(self, message):
    for product in self.products:
      position = ProductPosition(symbol=product.symbol,
                                 long_position=0,
                                 short_position=0,
                                 net_position=0)
      self.set_position(product, 0.0, position)

    msg = DeribitFuturesPrivateParser.parse_account_position(message)
    for position in msg.each_position:
      product = DeribitFuturesProduct.FromStrProduct(position.symbol)
      if product in self.products:
        self.set_position(product, position.net_position, position)

    self._unrealized_pnl = collections.defaultdict(float)
    self._avg_entry_price = collections.defaultdict(create_futures_positions)
    for msg in message['result']:
      direction = DeribitFuturesDirectionType(msg['direction'])
      product = DeribitFuturesProduct.FromStrNativeProduct(msg['instrument_name'])
      currency = product.base.currency
      self._unrealized_pnl[currency] += float(msg['floating_profit_loss'])
      if direction == DeribitFuturesDirectionType.BUY:
        self._avg_entry_price[product].long = float(msg['size'])
        self._avg_entry_price[product].long_avg_price = float(msg['average_price'])
      elif direction == DeribitFuturesDirectionType.SELL:
        self._avg_entry_price[product].short = -float(msg['size'])
        self._avg_entry_price[product].short_avg_price = float(msg['average_price'])
      elif direction == DeribitFuturesDirectionType.ZERO:
        self._avg_entry_price[product].short = float(msg['size'])
        self._avg_entry_price[product].long = float(msg['size'])
        self._avg_entry_price[product].long_avg_price = float(msg['average_price'])
        self._avg_entry_price[product].short_avg_price = float(msg['average_price'])

  def balance_as_proto(self):
    return AccountBalance(exchange='Deribit', each_balance=self.get_notNone_balances_internal())

  def position_as_proto(self):
    return AccountPosition(exchange='Deribit', each_position=self.get_notNone_positions_internal())

  def get_product_position_proto(self, product):
    return self._positions_internal[product.symbol]

  def get_futures_position(self, product):
    """return FuturesPosition namedtuple"""
    return self._avg_entry_price.get(product, None)

  # Margin balance includes unrealized pnl, but this is provided by exchange
  # and could be out of date or slow.
  # DO NOT USE THIS IN TRADING, ONLY ACCEPTABLE IN PNL DISPLAY.
  def get_margin_balance(self, currency: str):
    return self._unrealized_pnl[currency] + self.get_balance(currency)


class DeribitFuturesOrderGateway(RestOrderGateway):
  _exchange = 'Deribit'
  _market_type = 'Futures'
  _api_version = None

  OGInfoType = DeribitFuturesOrderGatewayInfo

  class _MsgType(enum.IntEnum):
    REST_ACCOUNT_BALANCE = 0
    REST_ACCOUNT_POSITION = 1
    OPEN_ORDERS = 2
    ALL_ORDERS = 3
    ORDERS_INFO = 4
    WS_ORDER_UPDATE = 5

  def __init__(self,
               currencies: typing.List[DeribitCurrency],
               products: typing.List[DeribitFuturesProduct],
               config: Config,
               og_config=None,
               logger=None):
    super().__init__(currencies, products, config, og_config, logger)
    self._og_ready = PositionOrderGatewayReady()
    self._og_ready.set_status(OrderGatewayStatus.INITIALIZING)

    # Following instance attributes are initialized in start method.
    self._flow_sub = None
    self._ioloop = None
    self._publisher = None
    self._scheduler = None
    self._private_client = None
    self._ws_private_client = None
    self._rest_rate_limiter = RestRateLimiter(delay=10)
    self._use_order_info = True

    self._og_config = OrderGatewayConfig(
        account_query_period=10,
        user_query_period=2,
        orders_snapshot_query_period=10,
        fill_query_period=2,
    )
    if og_config is not None:
      self._og_config.MergeFrom(og_config)
    self._og_logger.exchange_type = coin_enum.Deribit

  def submittable(self):
    return self.is_ready() and self._rest_rate_limiter.is_available()

  def get_deribit_futures_position(self, product: DeribitFuturesProduct):
    return self._og_info.get_futures_position(product)

  def start(self, queue, ioloop, flow_sub=None):
    api_host = self._og_config.rest_root_url if \
               self._og_config.HasField('rest_root_url') else None
    disable_ssl = self._og_config.rest_disable_ssl

    topic_id_str = 'deribit_futures_rest_private_1'
    topic_id = topic_map.search_topic_id(topic_id_str)[0][1]
    self._ioloop = ioloop

    self._publisher = queue.get_writer(topic_id=topic_id)
    self._scheduler = Scheduler(ioloop=ioloop, logger=self._logger)

    self._private_client = DeribitFuturesNativePrivateClient(key_file=self._config.key_file,
                                                             api_host=api_host,
                                                             disable_ssl=disable_ssl,
                                                             use_async=True)
    self._prepare_ws_private_client(start=True)
    self._install_periodic_query()

    # Subscribe to flow for feed.
    self._flow_sub = flow_sub or flow_subscriber.from_queue(queue)
    self._flow_sub.subscribe(topic_id_str, self.on_private_exchange_msg)

  @tornado.gen.coroutine
  def _install_periodic_query(self):
    assert self._scheduler is not None, 'self._scheduler is not set!'
    assert self._publisher is not None, 'self._publisher is not set!'
    symbols = [prod.base.native_currency for prod in self._og_info.products]

    for symbol in symbols:
      query_balance_task = wrap_async_task(self._private_client.query_account_summary,
                                           msg_type=self._MsgType.REST_ACCOUNT_BALANCE,
                                           publisher=self._publisher,
                                           func_params=dict(currency=symbol))
      # query balance
      self._scheduler.add_periodic_task(
          period=self._og_config.account_query_period,
          job_func=query_balance_task,
      )

    yield tornado.gen.sleep(self._og_config.burst_query_avoid_time)
    for symbol in symbols:
      query_position_task = wrap_async_task(self._private_client.query_positions,
                                            msg_type=self._MsgType.REST_ACCOUNT_POSITION,
                                            publisher=self._publisher,
                                            func_params=dict(currency=symbol))

      # query positions
      self._scheduler.add_periodic_task(
          period=self._og_config.user_query_period,
          job_func=query_position_task,
      )

    # query open orders
    for symbol in symbols:
      query_orders_task = wrap_async_task(self._private_client.query_open_orders,
                                          msg_type=self._MsgType.OPEN_ORDERS,
                                          publisher=self._publisher,
                                          func_params=dict(currency=symbol))
      self._scheduler.add_periodic_task(period=self._og_config.orders_snapshot_query_period,
                                        job_func=query_orders_task)

    yield tornado.gen.sleep(self._og_config.burst_query_avoid_time)
    for symbol in symbols:
      query_orders_task = wrap_async_task(self._private_client.query_historical_orders,
                                          msg_type=self._MsgType.ALL_ORDERS,
                                          publisher=self._publisher,
                                          func_params=dict(currency=symbol, count=30))
      self._scheduler.add_periodic_task(
          period=self._og_config.fill_query_period,
          job_func=query_orders_task,
      )

    def print_info():
      print_working_orders(self._logger, self)

    self._scheduler.add_periodic_task(
        period=10,
        job_func=print_info,
    )
    self._scheduler.start()

  def _prepare_ws_private_client(self, start=False):
    if self._ws_private_client is not None:
      self._ws_private_client.stop()

    self._ws_private_client = DeribitFuturesWsPrivateClient(
        self._config.key_file,
        ioloop=self._ioloop,
        notification_callback=self._notification_callback,
        close_callback=functools.partial(self._prepare_ws_private_client, True),
        on_ready=self._on_ws_private_client_ready,
    )

    if start:
      self._ws_private_client.start()

  def _on_ws_private_client_ready(self):
    native_symbols = [prod.native_symbol for prod in self._og_info.products]
    self._ws_private_client.prepare_subscription(native_symbols)

  def _notification_callback(self, data):
    assert data['method'] == 'subscription', data
    timestamp = get_timestamp()
    topic = data['params']['channel']
    if topic.startswith('user.orders.'):
      msg_type = self._MsgType.WS_ORDER_UPDATE
    else:
      self._logger.error('Unknown subscription topic! %s', data)
      return

    exchange_message = PrivateExchangeMessage(msg_type=msg_type,
                                              data=data,
                                              received_timestamp=timestamp,
                                              request_params=None,
                                              request_timestamp=None)
    self._publisher.write(timestamp, exchange_message)

  def _async_submit_order(self, order, post_only):
    assert order.order_id is not None, order
    assert order.order_type == OrderType.LIMIT, order

    formatter = self._order_formatter[order.product]
    # direction, offset = translate_order_side(order.order_side)
    params = {
        'symbol': order.product.native_symbol,
        'price': formatter.format_price(order.price),
        'qty': formatter.format_qty(order.qty),
        'side': DeribitFuturesDirectionType(order.order_side.name.lower()),
        'order_type': DeribitFuturesOrderPriceType(order.order_type.name.lower()),
        'post_only': post_only
    }
    fut = self._private_client.create_order(**params)
    return fut

  def _async_cancel_order(self, order):
    assert order.internal.exchange_order_id is not None, order
    real_exchange_order_id = get_real_exchange_order_id(order)
    fut = self._private_client.cancel_order(order_id=real_exchange_order_id,)
    return fut

  def _handle_submit_response(self, order, response):
    if response.status_code != requests.codes.ok:
      self._order_manager.remove_order(order, ignore_error=True)
      self._logger.error('Submit order %s failed!\n%s', order, response.content)
      return OrderEvent.ORDER_REJECTED

    message = response.json()
    if 'trades' not in message['result'] or 'order' not in message['result']:
      self._logger.error('Status is not ok! %s', message)
      self._order_manager.remove_order(order, ignore_error=True)
      return OrderEvent.ORDER_ERROR

    exchange_order_id = convert_to_exchange_order_id(order.product,
                                                     message['result']['order']['order_id'])
    prev_order = self._order_manager.get(exchange_order_id=exchange_order_id)
    if prev_order is None:
      self._order_manager.set_order(order, exchange_order_id=exchange_order_id)
    else:
      if prev_order.order_id != order.order_id:
        self._logger.info('order id replaced: %s => %s', prev_order.order_id, order.order_id)
        self._order_manager.remove_order(prev_order)
        self._order_manager.set_order(order, exchange_order_id=exchange_order_id)
        self._order_manager.remove_order(order)
        prev_order.order_id = order.order_id
        self._order_manager.add_order(prev_order)
    return OrderEvent.ORDER_ACCEPTED

  def _handle_cancel_response(self, order, response):
    if response.status_code != requests.codes.ok:
      self._logger.error('Cancel order failed with http response code=%s!. %s' %
                         (response.status_code, response.content))
      return OrderEvent.CANCEL_ERROR

    message = response.json()
    if 'result' not in message:
      self._logger.error('Cancel order (%s) error! %s', order.order_id, message)
      if self._order_manager.is_dead_order(order):
        self._logger.warning('Order (%s) is dead already!', order.order_id)
      else:
        self._logger.error('Unexpected error! Need inspection! %s', message)
      return OrderEvent.CANCEL_ERROR

    if not self._order_manager.is_dead_order(order):
      self._order_manager.set_order(order, cancel_confirmed=True)
    else:
      self._logger.info('Order is dead already. %s', order.order_id)
    return OrderEvent.CANCEL_CONFIRMED

  def _handle_error_response(self, response):
    if response.status_code == 504:
      self._logger.error('504 Gateway Timeout error!')
    else:
      self._logger.error('response error! %s, %s, %s',
                         response.status_code,
                         response.content,
                         response.url)

  @exchange_message_handler(_MsgType.OPEN_ORDERS)
  def _handle_open_orders(self, exchange_message: PrivateExchangeMessage):
    timestamp = get_timestamp()
    response = exchange_message.data
    if response.status_code != requests.codes.ok:
      self._handle_error_response(response)
      return

    message = response.json()
    if 'result' not in message:
      self._logger.error('status is not ok: %s' % message)
      return

    exchange_open_order_dict = {}
    products = set()
    for entry in message['result']:
      native_symbol = entry['instrument_name']
      product = DeribitFuturesProduct.FromStrNativeProduct(native_symbol)
      if product not in self._og_info.products:
        continue  # ignore

      order = create_order_from_open_order_message(entry, timestamp)
      exchange_order_id = order.internal.exchange_order_id
      prev_order = self._order_manager.get(exchange_order_id=exchange_order_id)
      if prev_order is None:
        order_id = self._order_id_generator.gen_order_id(is_foreign=True)
        order.order_id = order_id
      else:
        order.order_id = prev_order.order_id
      exchange_open_order_dict[exchange_order_id] = order

    self._order_manager.copy_from_exchange_open_orders(exchange_open_order_dict, products)

    if not self.is_ready():
      self._og_ready.set_orders_ready(True)
      self._og_ready.check_ready()

  @exchange_message_handler(_MsgType.ORDERS_INFO)
  def _handle_orders_info(self, exchange_message: PrivateExchangeMessage):
    timestamp = get_timestamp()
    response = exchange_message.data
    if response is None:
      return

    if response.status_code != requests.codes.ok:
      self._handle_error_response(response)
      return

    message = response.json()
    if 'result' not in message:
      self._logger.error('status is not ok: %s' % message)
      return

    self._handle_raw_order_info(message['result'], timestamp)

  @exchange_message_handler(_MsgType.ALL_ORDERS)
  def _handle_all_orders(self, exchange_message: PrivateExchangeMessage):
    timestamp = get_timestamp()
    response = exchange_message.data
    if response is None:
      return

    if response.status_code != requests.codes.ok:
      self._handle_error_response(response)
      return

    message = response.json()
    if 'result' not in message:
      self._logger.error('status is not ok: %s' % message)
      return

    for entry in message['result']:
      self._handle_raw_order_info(entry, timestamp)

  @exchange_message_handler(_MsgType.WS_ORDER_UPDATE)
  def _handle_ws_orders_update(self, exchange_message: PrivateExchangeMessage):
    timestamp = exchange_message.received_timestamp
    data = exchange_message.data
    self._handle_raw_order_info(data['params']['data'], timestamp)

  def _handle_raw_order_info(self, raw_order_info, timestamp):
    native_symbol = raw_order_info['instrument_name']
    product = DeribitFuturesProduct.FromStrNativeProduct(native_symbol)
    if product not in self._og_info.products:
      return  # ignore

    order_info = create_order_from_open_order_message(raw_order_info, timestamp)
    exchange_order_id = order_info.internal.exchange_order_id
    order = self._order_manager.get(exchange_order_id=exchange_order_id)
    if order is None:
      order_id = self._order_id_generator.gen_order_id(is_foreign=True)
      order_info.order_id = order_id
      self._order_manager.add_order(order_info)
      order = order_info
    else:
      order_info.order_id = order.order_id
      order.internal.fee = order_info.internal.fee
    self._update_order(order, order_info, timestamp)

  def _update_order(self, order, order_info, timestamp):
    status = order_info.internal.status
    exchange_order_id = order_info.internal.exchange_order_id
    if order.accepted is False:
      self._order_manager.set_order(order, exchange_order_id=exchange_order_id)

    if status == DeribitFuturesStatusType.OPEN:
      pass
    elif status == DeribitFuturesStatusType.FILLED:
      if not self._order_manager.is_dead_order(order) and order_info.internal.fill.fully_filled:
        self._order_manager.set_order(order, fully_filled=True)
    elif status in (DeribitFuturesStatusType.CANCELLED, DeribitFuturesStatusType.REJECTED):
      if not self._order_manager.is_dead_order(order):
        self._order_manager.set_order(order, cancel_confirmed=True)
    else:
      self._logger.info('%s(%s) status: %s',
                        order.order_id,
                        order.internal.exchange_order_id,
                        status)

    price_avg_fill = order_info.internal.fill.value.price_avg_fill
    qty_accum_fill = order_info.internal.fill.value.qty_accum_fill
    last_fill = order.internal.fill.update_by_accum_fill(price_avg_fill, qty_accum_fill)
    if last_fill:
      self._publish_fill(last_fill.price_last_fill, last_fill.qty_last_fill, order)
      self._og_logger.gen_order_filled_event_and_log(order, timestamp)
      order.internal.fill.log_last_fill(order, self._logger)
    elif order is order_info and order.internal.status == DeribitFuturesStatusType.FILLED:
      price_last_fill = order.internal.fill.value.price_last_fill
      qty_last_fill = order.internal.fill.value.qty_last_fill
      self._publish_fill(price_last_fill, qty_last_fill, order)
      self._og_logger.gen_order_filled_event_and_log(order, timestamp)
      order.internal.fill.log_last_fill(order, self._logger)

  @exchange_message_handler(_MsgType.REST_ACCOUNT_BALANCE)
  def _handle_rest_balance(self, exchange_message: PrivateExchangeMessage):
    response = exchange_message.data
    if response.status_code != requests.codes.ok:
      self._handle_error_response(response)
      return

    message = response.json()
    if 'result' not in message:
      self._logger.error('status is not ok: %s' % message)
      return

    self._og_info.update_balance(message)
    self._og_logger.write_balance(self._og_info.balance_as_proto())
    if not self.is_ready():
      self._og_ready.set_balance_ready(True)
      self._og_ready.check_ready()

  @exchange_message_handler(_MsgType.REST_ACCOUNT_POSITION)
  def _handle_rest_position(self, exchange_message: PrivateExchangeMessage):
    response = exchange_message.data
    if response.status_code != requests.codes.ok:
      self._handle_error_response(response)
      return

    message = response.json()
    if 'result' not in message:
      self._logger.error('status is not ok: %s' % message)
      return

    self._og_info.update_position(message)
    self._og_logger.write_position(self._og_info.position_as_proto())
    if not self.is_ready():
      self._og_ready.set_position_ready(True)
      self._og_ready.check_ready()
